package global

import (
	"encoding/json"
	"fmt"
	"gitee.com/jamon/espoir_go/common"
	"gitee.com/jamon/espoir_go/common/utils"
	"gitee.com/jamon/espoir_go/logger"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/xtaci/kcp-go"
	"sync"
)

type SessionInfo struct {
	sessionID string
	token     string
	protoType uint8
}

type SessionManager struct {
	sync.Mutex
	lastSessionSeq uint32

	// session_id: sessionInfo
	sessionMap utils.Map[string, *SessionInfo]

	// session_id: web_conn
	websocketConnMap utils.Map[string, *websocket.Conn]

	// session_id: kcp_conn
	kcpConnMap utils.Map[string, *kcp.UDPSession]
}

var (
	sessionManager *SessionManager
	sessionOnce    sync.Once
)

func GetSessionManager() *SessionManager {
	sessionOnce.Do(func() {
		sessionManager = &SessionManager{
			lastSessionSeq:   0,
			sessionMap:       utils.NewMap[string, *SessionInfo](),
			websocketConnMap: utils.NewMap[string, *websocket.Conn](),
			kcpConnMap:       utils.NewMap[string, *kcp.UDPSession](),
		}
	})
	return sessionManager
}

func (mgr *SessionManager) getNextSessionSeq() uint32 {
	mgr.Lock()
	defer mgr.Unlock()
	mgr.lastSessionSeq++
	return mgr.lastSessionSeq
}

func (mgr *SessionManager) GetNextSessionId() (sessionID string, err error) {
	sessionID = fmt.Sprintf("%s_%d", GetNodeData().Config.NodeID, mgr.getNextSessionSeq())
	return
}

func (mgr *SessionManager) GetToken() (token string) {
	rawToken, err := uuid.NewUUID()
	if err != nil {
		logger.Error("GetToken: ", logger.Kwargs{"err": err})
		return ""
	}
	return rawToken.String()
}

func (mgr *SessionManager) addSession(sessionID string, protoType uint8) {
	token := mgr.GetToken()
	mgr.sessionMap.Put(sessionID, &SessionInfo{
		sessionID: sessionID,
		token:     token,
		protoType: protoType,
	})
}

func (mgr *SessionManager) IsSessionValid(sessionID string, token string) bool {
	rawToken, ok := mgr.sessionMap.Get(sessionID)
	if !ok {
		return false
	}
	return rawToken.token == token
}

func (mgr *SessionManager) AddWebsocketSession(conn *websocket.Conn) (sessionID string, err error) {
	sessionID, err = mgr.GetNextSessionId()
	if err != nil {
		logger.Error("AddWebsocketSession: ", logger.Kwargs{"err": err})
		return
	}
	mgr.addSession(sessionID, common.PROTO_WEBSOCKET)
	mgr.websocketConnMap.Put(sessionID, conn)
	return
}

func (mgr *SessionManager) AddKcpSession(conn *kcp.UDPSession) (sessionID string, err error) {
	sessionID, err = mgr.GetNextSessionId()
	if err != nil {
		logger.Error("AddKcpSession: ", logger.Kwargs{"err": err})
		return
	}
	mgr.addSession(sessionID, common.PROTO_KCP)
	mgr.kcpConnMap.Put(sessionID, conn)
	return
}

func (mgr *SessionManager) RemoveSession(sessionID string) {
	sessionInfo, ok := mgr.sessionMap.Get(sessionID)
	if !ok {
		return
	}
	pType := sessionInfo.protoType
	if pType == common.PROTO_WEBSOCKET {
		mgr.websocketConnMap.Remove(sessionID)
	} else if pType == common.PROTO_KCP {
		mgr.kcpConnMap.Remove(sessionID)
	}
	mgr.sessionMap.Remove(sessionID)
}

func (mgr *SessionManager) GetWebsocketConn(sessionID string) (conn *websocket.Conn, err error) {
	conn, ok := mgr.websocketConnMap.Get(sessionID)
	if !ok {
		err = fmt.Errorf("no conn for session_id: %v", sessionID)
		return
	}
	return
}

func (mgr *SessionManager) UpdateWebsocketConn(sessionID string, conn *websocket.Conn) {
	mgr.websocketConnMap.Put(sessionID, conn)
}

func (mgr *SessionManager) SendMsg(sessionID string, res *common.MsgResponse) (err error) {
	conn, err := mgr.GetWebsocketConn(sessionID)
	if err != nil {
		return
	}
	msg, err := json.Marshal(res)
	if err != nil {
		logger.Error("SendMsg: ", logger.Kwargs{"err": err})
		return
	}
	logger.Debug("SendMsg: ", logger.Kwargs{"msg": string(msg)})
	err = conn.WriteMessage(websocket.TextMessage, msg)
	if err != nil {
		logger.Error("SendMsg: ", logger.Kwargs{"err": err})
		return
	}
	return
}

func (mgr *SessionManager) Broadcast(res *common.MsgResponse) (err error) {
	mgr.websocketConnMap.LockRange(func(key string, value *websocket.Conn) {
		go common.SafeRun(func() {
			err = mgr.SendMsg(key, res)
			if err != nil {
				logger.Error("Broadcast: ", logger.Kwargs{"err": err})
				return
			}
			return
		})
	})
	return
}
